import os
import requests
import os
import json
import time
from datetime import datetime

originPath = "/home/cyw/projects/data/functionsim/allorigindata"
finishedName = "/home/cyw/projects/get_bianry_sample_famLable/finishedName.txt"
VTJsonSaveFile = "/home/cyw/projects/get_bianry_sample_famLable/lableJson.txt"
vtLogPath = "/home/cyw/projects/get_bianry_sample_famLable/vtlogs.txt"


def write_file(path, str1):
    with open(path, "a+") as f:
        f.write(str(datetime.now())+"\t"+str1+"\n")


def get_sample_names():
    """
        获得需要处理的样本
    """
    global originPath, finishedName
    tar = {}
    m, n = 0, 0
    g = os.walk(originPath)
    for path, dir_list, file_list in g:
        for file_name in file_list:
            tar[file_name] = os.path.join(path, file_name)
    m = len(tar)

    with open(finishedName) as file:
        for i in file.readlines():
            tempName = i.strip()
            if tempName in tar:
                del tar[tempName]
                # print("--")
    n = len(tar)
    print("共{}个样本，还需处理{}个样本".format(m, n))
    return tar


def get_json_from_VT(name, visitSpan):
    """
        将样本上传到vt获得json报告
    """
    global vtLogPath, VTJsonSaveFile, finishedName
    url = "https://www.virustotal.com/api/v3/files"
    headers = {
        "accept": "application/json",
        "x-apikey": "b619a388252fb498b88e56e2efec3c0aa70be5edbb891494c5b14c4deb4e582c"
    }

    files = {"file": open(name, "rb")}
    write_file(vtLogPath, str(files))
    response = requests.post(url, files=files, headers=headers)
    write_file(vtLogPath, response.text)
    data = json.loads(response.text)
    analysisId = data["data"]["id"]
    write_file(vtLogPath, name+"---->analysis id = "+analysisId)
    # 查看分析情况，completed为分析完成

    # analyses不消耗quota
    url1 = "https://www.virustotal.com/api/v3/analyses/"+analysisId
    write_file(vtLogPath, "等待virusTotal处理样本")
    time.sleep(visitSpan)
    while (True):
        response = requests.get(url1, headers=headers)
        data = json.loads(response.text)
        sha256 = data["meta"]["file_info"]["sha256"]
        is_avaliable = True if data["data"]["attributes"]["status"] == "completed" else False
        if (is_avaliable):
            write_file(vtLogPath, "样本分析完成")
            break
        write_file(vtLogPath, "---")
        time.sleep(visitSpan)

    # 获取json报告
    url2 = "https://www.virustotal.com/api/v3/files/"+sha256
    response = requests.get(url2, headers=headers)
    data = json.loads(response.text)
    temp = json.dumps(data)

    with open(VTJsonSaveFile, "a+") as f:
        f.write(temp)
        f.write("\n")
    with open(finishedName, "a+") as f:
        f.write(name.split("/")[-1]+"\n")


def analyse_vtJson_by_avclass():
    # python3 /home/cyw/projects/get_bianry_sample_famLable/avclass/avclass/avclass_labeler.py -vt /home/chenyongwei/data/functionsim/predata/unLableJson.txt -vt3 -c
    pass


#   nohup python3 /home/cyw/projects/get_bianry_sample_famLable/getVTJsonReport.py &
if __name__ == "__main__":
    # tar=get_sample_names()
    # cnt=1
    # for name in tar.keys():
    #     write_file(vtLogPath,"--------start analyse---------")
    #     write_file(vtLogPath,"第{}个样本进行分析".format(cnt))
    #     print("第{}个样本进行分析".format(cnt))
    #     # 350的时间间隔差不多能跑满
    #     # get_json_from_VT(tar[name],200)
    #     get_json_from_VT(tar[name],320)
    #     cnt+=1

    # 这里进行少量样本的补充
    print("开始处理")
    # name=["12476aa0d72285be95f0c8c125a0bd36","6dcb616a8545920c8b128de2359c6f22" ,"9ff91b23074ae660164df7f0b2faddea","2a8ea48b988d141fc33bb17ea5b0979e"]
    name = ['022129753556a19e4a85c4ad8fbfa785', '742a22390bc72ec38eb18c95446972e8', '8ee6ae272842237d5fa3f7a2347631bb', 'e9ebe89830cefa1182c50a3821f69d23', 'b8b4c27058b737bc4149879395bbfa0a', '78d7fff3916ce40dfe50a50277ff60f8', 'ca4812960ec2e217cb8b43f80ba8447b',
            'ccf428c49052e481fd9e8971eae948ce', 'f928a7e68d826fce5cc4e87e943341d7', 'bf313223cd0c1bdf938c22eb7344371f', 'fd7eb7aa093bd5691214caa046b9b890', '3fe4ce7e3d6d06f6894b23c03cfcb46a', 'c8e1f8ecb5ba5a0a3c9338c0715abc35', 'df0b77aa2a00b8fe204e91baa8a3bed4']
    for i in name:
        get_json_from_VT(originPath+"/"+i, 30)
        print("--")
